#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
线性回归算法演示项目
基于CSDN博客：https://blog.csdn.net/weixin_50804299/article/details/136328878
实现了糖尿病数据预测和波士顿房价预测两个案例
"""

import numpy as np
import pandas as pd
import matplotlib

matplotlib.use('Agg')  # 使用非交互式后端，避免GUI问题
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.datasets import load_diabetes
import warnings

warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei']
plt.rcParams['axes.unicode_minus'] = False


class LinearRegressionDemo:
    def __init__(self):
        self.results = {}

    def diabetes_prediction(self):
        """糖尿病数据线性回归预测"""
        print("=" * 50)
        print("糖尿病数据线性回归预测")
        print("=" * 50)

        # 1. 加载数据
        diabetes = load_diabetes()
        data = diabetes['data']
        target = diabetes['target']
        feature_names = diabetes['feature_names']

        # 2. 创建DataFrame
        df = pd.DataFrame(data, columns=feature_names)

        # 3. 划分训练集和测试集
        x_train, x_test, y_train, y_test = train_test_split(
            data, target, test_size=0.2, random_state=42
        )

        # 4. 创建并训练模型
        model = LinearRegression()
        model.fit(x_train, y_train)

        # 5. 预测
        y_pred = model.predict(x_test)

        # 6. 评估指标
        score = model.score(x_test, y_test)
        mae = mean_absolute_error(y_test, y_pred)
        mse = mean_squared_error(y_test, y_pred)
        rmse = np.sqrt(mse)

        print(f"模型得分 (R²): {score:.4f}")
        print(f"平均绝对误差 (MAE): {mae:.4f}")
        print(f"均方误差 (MSE): {mse:.4f}")
        print(f"均方根误差 (RMSE): {rmse:.4f}")

        # 7. 特征重要性分析
        print("\n特征系数 (影响程度):")
        for name, coef in zip(feature_names, model.coef_):
            print(f"{name}: {coef:.4f}")

        # 8. 绘制特征与目标关系图
        plt.figure(figsize=(20, 10))
        for i, col in enumerate(df.columns):
            ax = plt.subplot(2, 5, i + 1)

            # 散点图
            ax.scatter(df[col], target, alpha=0.5)

            # 拟合直线
            single_feature_model = LinearRegression()
            single_feature_model.fit(df[[col]], target)
            x_line = np.array([df[col].min(), df[col].max()])
            y_line = single_feature_model.predict(x_line.reshape(-1, 1))
            ax.plot(x_line, y_line, 'r-', linewidth=2)

            ax.set_title(f'{col}\n系数: {single_feature_model.coef_[0]:.3f}')
            ax.set_xlabel(col)
            ax.set_ylabel('Target')

        plt.tight_layout()
        plt.savefig('diabetes_features_analysis.png', dpi=300, bbox_inches='tight')
        plt.close()

        # 9. 保存结果
        self.results['diabetes'] = {
            'score': score,
            'mae': mae,
            'mse': mse,
            'rmse': rmse,
            'coefficients': dict(zip(feature_names, model.coef_)),
            'intercept': model.intercept_
        }

        return self.results['diabetes']

    def boston_housing_prediction(self):
        """波士顿房价预测"""
        print("\n" + "=" * 50)
        print("波士顿房价预测")
        print("=" * 50)

        try:
            # 由于load_boston在新版本中已弃用，使用替代数据
            # 创建模拟的波士顿房价数据
            np.random.seed(42)
            n_samples = 506

            # 生成特征数据
            crim = np.random.exponential(3.61, n_samples)
            zn = np.random.exponential(11.36, n_samples)
            indus = np.random.normal(11.14, 6.86, n_samples)
            chas = np.random.choice([0, 1], n_samples, p=[0.93, 0.07])
            nox = np.random.normal(0.55, 0.12, n_samples)
            rm = np.random.normal(6.28, 0.70, n_samples)
            age = np.random.exponential(68.57, n_samples)
            dis = np.random.exponential(3.80, n_samples)
            rad = np.random.choice(range(1, 25), n_samples)
            tax = np.random.normal(408.24, 168.54, n_samples)
            ptratio = np.random.normal(18.46, 2.16, n_samples)
            b = np.random.normal(356.67, 91.29, n_samples)
            lstat = np.random.exponential(12.65, n_samples)

            # 生成目标变量（房价）
            price = (
                    -0.1 * crim + 0.05 * zn - 0.8 * indus + 2.5 * chas - 30 * nox +
                    5 * rm - 0.05 * age - 1.5 * dis - 0.5 * rad - 0.01 * tax -
                    1.2 * ptratio + 0.01 * b - 0.8 * lstat + 25 + np.random.normal(0, 5, n_samples)
            )

            # 创建DataFrame
            boston_data = pd.DataFrame({
                'CRIM': crim, 'ZN': zn, 'INDUS': indus, 'CHAS': chas,
                'NOX': nox, 'RM': rm, 'AGE': age, 'DIS': dis,
                'RAD': rad, 'TAX': tax, 'PTRATIO': ptratio, 'B': b,
                'LSTAT': lstat, 'PRICE': price
            })

            # 数据探索
            print("数据基本信息:")
            print(boston_data.info())
            print("\n缺失值统计:")
            print(boston_data.isnull().sum())
            print("\n描述性统计:")
            print(boston_data.describe())

            # 相关性分析
            corr_matrix = boston_data.corr()
            plt.figure(figsize=(12, 10))
            sns.heatmap(corr_matrix, annot=True, cmap='RdGy', center=0)
            plt.title('波士顿房价数据相关性热力图')
            plt.tight_layout()
            plt.savefig('boston_correlation.png', dpi=300, bbox_inches='tight')
            plt.close()

            # 查尔斯河影响分析
            chas_analysis = boston_data.groupby('CHAS')['PRICE'].agg(['mean', 'count'])
            print("\n查尔斯河对房价的影响:")
            print(chas_analysis)

            # 可视化查尔斯河影响
            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

            # 饼图
            chas_counts = boston_data['CHAS'].value_counts()
            ax1.pie(chas_counts.values, labels=['未穿过河流', '穿过河流'],
                    autopct='%1.1f%%', startangle=90)
            ax1.set_title('查尔斯河分布')

            # 箱线图
            boston_data.boxplot(column='PRICE', by='CHAS', ax=ax2)
            ax2.set_title('查尔斯河对房价的影响')
            ax2.set_xlabel('是否穿过查尔斯河')
            ax2.set_ylabel('房价')

            plt.tight_layout()
            plt.savefig('boston_chas_analysis.png', dpi=300, bbox_inches='tight')
            plt.close()

            # 特征与房价散点图
            features = ['CRIM', 'ZN', 'INDUS', 'CHAS', 'NOX', 'RM',
                        'AGE', 'DIS', 'RAD', 'TAX', 'PTRATIO', 'B', 'LSTAT']

            plt.figure(figsize=(20, 15))
            for i, feature in enumerate(features):
                plt.subplot(4, 4, i + 1)
                plt.scatter(boston_data[feature], boston_data['PRICE'], alpha=0.5, s=10)
                plt.xlabel(feature)
                plt.ylabel('PRICE')
                plt.title(f'{feature} vs PRICE')

            plt.tight_layout()
            plt.savefig('boston_scatter_plots.png', dpi=300, bbox_inches='tight')
            plt.close()

            # 线性回归建模
            X = boston_data[features]
            y = boston_data['PRICE']

            # 划分训练集和测试集
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=0.2, random_state=42
            )

            # 创建并训练模型
            model = LinearRegression()
            model.fit(X_train, y_train)

            # 预测
            y_pred = model.predict(X_test)

            # 评估
            score = r2_score(y_test, y_pred)
            mae = mean_absolute_error(y_test, y_pred)
            mse = mean_squared_error(y_test, y_pred)
            rmse = np.sqrt(mse)

            print(f"模型得分 (R²): {score:.4f}")
            print(f"平均绝对误差 (MAE): {mae:.4f}")
            print(f"均方误差 (MSE): {mse:.4f}")
            print(f"均方根误差 (RMSE): {rmse:.4f}")

            # 回归系数
            print("\n特征回归系数:")
            for feature, coef in zip(features, model.coef_):
                print(f"{feature}: {coef:.4f}")
            print(f"截距: {model.intercept_:.4f}")

            # 预测结果可视化
            plt.figure(figsize=(12, 5))

            # 实际值vs预测值
            plt.subplot(1, 2, 1)
            plt.scatter(y_test, y_pred, alpha=0.5)
            plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
            plt.xlabel('实际房价')
            plt.ylabel('预测房价')
            plt.title('实际值 vs 预测值')

            # 残差图
            plt.subplot(1, 2, 2)
            residuals = y_test - y_pred
            plt.scatter(y_pred, residuals, alpha=0.5)
            plt.axhline(y=0, color='r', linestyle='--')
            plt.xlabel('预测房价')
            plt.ylabel('残差')
            plt.title('残差分析')

            plt.tight_layout()
            plt.savefig('boston_prediction_results.png', dpi=300, bbox_inches='tight')
            plt.close()

            # 保存结果
            self.results['boston'] = {
                'score': score,
                'mae': mae,
                'mse': mse,
                'rmse': rmse,
                'coefficients': dict(zip(features, model.coef_)),
                'intercept': model.intercept_
            }

            return self.results['boston']

        except Exception as e:
            print(f"波士顿房价预测出错: {e}")
            return None

    def run_all(self):
        """运行所有演示"""
        print("开始线性回归算法演示...")

        # 运行糖尿病预测
        diabetes_result = self.diabetes_prediction()

        # 运行波士顿房价预测
        boston_result = self.boston_housing_prediction()

        # 总结报告
        print("\n" + "=" * 50)
        print("演示总结")
        print("=" * 50)

        if diabetes_result:
            print("糖尿病预测:")
            print(f"  R²得分: {diabetes_result['score']:.4f}")
            print(f"  RMSE: {diabetes_result['rmse']:.4f}")

        if boston_result:
            print("\n波士顿房价预测:")
            print(f"  R²得分: {boston_result['score']:.4f}")
            print(f"  RMSE: {boston_result['rmse']:.4f}")

        print("\n所有图表已保存为PNG文件:")
        print("1. diabetes_features_analysis.png - 糖尿病特征分析")
        print("2. boston_correlation.png - 波士顿房价相关性")
        print("3. boston_chas_analysis.png - 查尔斯河影响分析")
        print("4. boston_scatter_plots.png - 特征散点图")
        print("5. boston_prediction_results.png - 预测结果分析")

        return self.results


if __name__ == "__main__":
    demo = LinearRegressionDemo()
    results = demo.run_all()